package 力扣._00周赛._426_周赛.问题3;

import java.util.*;

class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int m = edges1.length + 1, n = edges2.length + 1;
        List<List<Integer>> adj = new ArrayList<>(m + n);
        for (int i = 0; i < m + n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges1) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }
        for (int[] edge : edges2) {
            adj.get(edge[0] + m).add(edge[1] + m);
            adj.get(edge[1] + m).add(edge[0] + m);
        }


        int[] res = new int[m];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                adj.get(i).add(j + m);
                adj.get(j + m).add(i);

                int count = bfs(adj, i, m + n, k);
                res[i] = Math.max(res[i], count);

                adj.get(i).remove(Integer.valueOf(j + m));
                adj.get(j + m).remove(Integer.valueOf(i));
            }
        }
        return res;
    }

    private int bfs(List<List<Integer>> adj, int start, int size, int k) {
        Queue<Integer> queue = new LinkedList<>();
        boolean[] visited = new boolean[size];
        queue.add(start);
        visited[start] = true;
        int level = 0, count = 0;

        while (!queue.isEmpty() && level < k) {
            int levelSize = queue.size();
            for (int i = 0; i < levelSize; i++) {
                int node = queue.poll();
                for (int neighbor : adj.get(node)) {
                    if (!visited[neighbor]) {
                        visited[neighbor] = true;
                        queue.add(neighbor);
                    }
                }
            }
            level++;
        }
        for (int i = 0; i < size; i++) {
            if (visited[i]) {
                count++;
            }
        }
        return count;
    }

    public static void main(String[] args) {
        int[][] edges1 = new int[][]{{0, 1}, {0, 2}, {2, 3}, {2, 4}};
        int[][] edges2 = new int[][]{{0, 1}, {0, 2}, {0, 3}, {2, 7}, {1, 4}, {4, 5}, {4, 6}};
        int k = 2;
        Solution solution = new Solution();
        int[] res = solution.maxTargetNodes(edges1, edges2, k);
        for (int i : res) {
            System.out.println(i);
        }
    }
}